#include "Sample.h"
#include "../../Basis/Vector2.h"
#include <cmath>

Sampler::Sampler(unsigned long long _seed /* = 7564231ULL */)
	: m_rng(_seed)
{

}

//	2D sampling
//	generate num_samples samples in [0, 1) * [0, 1) for a single pixel

//	randomly generate num_samples
void Sampler::random(Vector2* samples, int num_samples)
{
	for (int i = 0; i < num_samples; i++)
	{
		samples[i].setX(m_rng());
		samples[i].setY(m_rng());
	}
}

//	assume num_samples is a perfect square, denote n = sqrt(num_samples)
//	divide the pixel into n * n blocks
//	generate a sample for each block
void Sampler::jitter(Vector2* samples, int num_samples)
{
	int n = (int)sqrt((float)num_samples);

	float x, y;
	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			x = (float)(i + m_rng()) / n;
			y = (float)(j + m_rng()) / n;
			samples[i * n + j].setX(x);
			samples[i * n + j].setY(y);
		}
	}
}

//	divide the pixel into num_samples * num_samples blocks
//	generate num_samples samples so that in each row
//	and in each column there is only one sample
void Sampler::nrooks(Vector2* samples, int num_samples)
{
	//	initialize all samples so that they distribute in the diagonal
	for (int i = 0; i < num_samples; i++)
	{
		samples[i].setX((float)(i + m_rng()) / num_samples);
		samples[i].setY((float)(i + m_rng()) / num_samples);
	}
	//	swap the e[0] value of all the samples
	//	it is a permutation in x-coordinates
	int index;
	float temp;
	for (int i = num_samples - 1; i >= 1; i--)
	{
		//	generate an index from 0, 1, 2, ..., i
		index = (int)(m_rng() * (i + 1));
		//	swap samples[i].e[0] and samples[index].e[0]
		temp = samples[i].x();
		samples[i].setX(samples[index].x());
		samples[index].setX(temp);
	}
}

//	combine jitter and nrooks together
//	assume num_samples is a perfect square
//	denote n as sqrt(num_samples)
//	generate num_samples samples which satisfy nrooks AND
//	divide the pixel into n * n blocks, in each blocks
//	there is only one sample

//	WARNING: THIS CODE IS REALLY REALLY HARD TO UNDERSTAND !!!
//	THE COMMENTS ARE ADDED BY ME
//	taodu@stanford.edu
//	TO BETTER UNDERSTAND THE CODE, WE WILL INTRODUCE AN EXAMPLE
//	WITH num_samples = 9
void Sampler::multiJitter(Vector2* samples, int num_samples)
{
	int n = (int)sqrt((float)num_samples);
	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			samples[i * n + j].setX((float)(i * n + j + m_rng()) / num_samples);
			samples[i * n + j].setY((float)(j * n + i + m_rng()) / num_samples);
		}
	}
	//	the result of the previous code will produce samples like this:
	//	--------------------------------------	X axis
	//	||0	|	|	||	|	|	||	|	|	||
	//	||	|	|	||3	|	|	||	|	|	||
	//	||	|	|	||	|	|	||6	|	|	||
	//	--------------------------------------
	//	||	|1	|	||	|	|	||	|	|	||
	//	||	|	|	||	|4	|	||	|	|	||
	//	||	|	|	||	|	|	||	|7	|	||
	//	--------------------------------------
	//	||	|	|2	||	|	|	||	|	|	||
	//	||	|	|	||	|	|5	||	|	|	||
	//	||	|	|	||	|	|	||	|	|8	||
	//	--------------------------------------
	//	Y axis

	
	int index;
	float temp;
	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			//	generate index from j, j + 1, ..., n - 1
			index = j + (int)(m_rng() * (n - j));
			
			//	swap the x coordinates of (i * n + j) and (i * n + index)
			temp = samples[i * n + j].x();
			samples[i * n + j].setX(samples[i * n + index].x());
			samples[i * n + index].setX(temp);

			index = j + (int)(m_rng() * (n - j));
			//	swap the y coordinates of (j * n + i) and (j * n + index)
			temp = samples[j * n + i].y();
			samples[j * n + i].setY(samples[j * n + index].y());
			samples[j * n + index].setY(temp);
		}
	}

	//	the for-loop above are REALLY unintuitive
	//	to better understand it, we follow these steps:
	
	//	first, the only thing it does is SWAP two samples' coordinates
	//	in which the two samples are selected from the same BIGGER row or column
	//	(separated by || or -- in the example above)
	//	which means: after each swap, the samples still keep the two properties:
	//	a. jitter AND b. nrook

	//	next, let's consider what the code has done in the swap
	//	the swap is the 2D version of the classic permutation problem
	//	consider when i = j = 0;
	//	the first swap selects a x position UNIFORMLY for sample 0
	//	the second swap selects a y position UNIFORMLY for sample 0
	
	//	then consider i = 0, j = 1;
	//	the first swap selects a x position UNIFORMLY for sample 1
	//	the second swap selects a y position UNIFORMLY for sample 3
	
	//	then i = 0, j = 2;
	//	the first swap selects a x position UNIFORMLY for sample 2
	//	the second swap selects a y position UNIFORMLY for sample 6
	//	etc..

	//	IMPORTANT:
	//	generally, in the ij loop,
	//	the first swap selects a x position UNIFORMLY for i * n + j
	//	the second swap selects a y position UNIFORMLY for j * n + i

	//	in the end, the samples distributed UNIFORMLY in the 2 dimensions
	//	and keep the two properties: jitter and nrook
}


//	wrap function transform the random number in [0, 1) into [-0.5, 0.5)
//	different wrap functions give different weight to the number
//	box filter give the same weight to all samples
void Sampler::boxFilter(Vector2* samples, int num_samples)
{
	for (int i = 0; i < num_samples; i++)
	{
		samples[i].setX(samples[i].x() - 0.5f);
		samples[i].setY(samples[i].y() - 0.5f);
	}
}

//	tent function transform [0, 1) into [-1, 1)
//	f(x) = sqrt(2x) - 1 (x < 0.5)
//	f(x) = 1 - sqrt(2-2x) (x > 0.5)
//	the samples are dense near zero
void Sampler::tentFilter(Vector2* samples, int num_samples)
{
	float x, y;
	for (int i = 0; i < num_samples; i++)
	{
		x = samples[i].x();
		y = samples[i].y();
		samples[i].setX(tentFunction(x));
		samples[i].setY(tentFunction(y));
	}
}

//	1D sampling
//	generate num_samples samples in [0, 1]

void Sampler::random(float* samples, int num_samples)
{
	for (int i = 0; i < num_samples; i++)
		samples[i] = m_rng();
}

//	divide [0, 1] into num_samples segments
//	so that each segment contains only one sample
void Sampler::jitter(float* samples, int num_samples)
{
	for (int i = 0; i < num_samples; i++)
		samples[i] = (float)(i + m_rng()) / num_samples;
}

float Sampler::tentFunction(float sample)
{
	if (sample < 0.5f)
		return sqrt(2.0f * sample) - 1.0f;
	else
		return 1.0f - sqrt(2.0f - 2.0f * sample);
}